from ultralytics import YOLO

def main():
    # 加载yolov8n预训练模型
    model = YOLO('yolov8n.pt')

    # 获取用户自定义的数据集 epochs迭代100次
    model.train(data="safehat.yaml", epochs=100)

    # 使用验证集验证效果
    model.val()

if __name__ == '__main__':
    main()